import subprocess
import time
from threading import Thread

from constant import *
from utils import *


class Env:
    def __init__(self):
        self.buffer_handler: Thread | None = None
        self.suricata_handler: subprocess.Popen[bytes] | None = None
        self.buffer = ""
        self.setup_env()

    def setup_suricata(self):
        self.suricata_handler = subprocess.Popen(SURICATA_SET_COMMAND, shell=True, cwd=WORK_DIR, stdout=subprocess.PIPE,
                                                 stderr=subprocess.PIPE)
        self.buffer_handler = Thread(target=self._output_buffer)
        self.buffer_handler.start()
        time.sleep(1)
        self.get_buffer()
        logger.success(f"Suricata setup successfully")

    def _output_buffer(self):
        while self.suricata_handler.poll() is None:
            line = self.suricata_handler.stderr.readline().decode("utf-8")
            self.buffer += line.strip()

    def get_buffer(self):
        ret = self.buffer
        self.buffer = ""
        return ret

    def terminate_suricata(self):
        if self.suricata_handler:
            self.suricata_handler.terminate()
        else:
            logger.warning("Suricata has not been setup yet.")

    def setup_env(self):
        if not is_sudo():
            logger.warning("ATTENTION! IPS agent is not running on sudo privilege. "
                           "Some functions may not running currently.")
        if not check_iptables():
            set_iptables()
        self.setup_suricata()

    # 必须要清理suricata
    # 不正确地停止有可能会导致NFQ队列出问题，如果除了问题，重启计算机或更换队列可解决
    def __del__(self):
        clear_iptables()
        self.terminate_suricata()
        logger.success("Suricata 正常退出")
        # clear_suricata_log()


def check_iptables() -> None | str:
    """
    Check if iptables is set
    :return: if iptables is set,return the lineno of NFQUEUE; or None if iptables is not set
    """
    ret = subprocess.run(IPTABLES_CHECK_COMMAND,
                         shell=True, cwd=WORK_DIR, stdout=subprocess.PIPE)
    if ret.stdout:
        return ret.stdout.decode().split(" ")[0]
    else:
        return None


def clear_iptables():
    if lineno := check_iptables():
        ret = subprocess.run(f"sudo iptables -D INPUT {lineno}", shell=True, stderr=subprocess.PIPE)
        if msg := ret.stderr:
            logger.error(f"Clear iptables error <sudo iptables -D DOCKER-USER {lineno}>, stderr:\n{msg}")


def set_iptables():
    ret = subprocess.run(IPTABLES_SET_COMMAND, shell=True, stderr=subprocess.PIPE)
    if msg := ret.stderr:
        logger.error(f"Set iptables error, stderr:\n{msg}")


def clear_suricata_log():
    with open(SURICATA_OUTPUT_FILE, "w") as f:
        f.write("")
    logger.info("Suricata log file cleared.")


env = Env()

if __name__ == "__main__":
    # time.sleep(10)
    clear_iptables()
